import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt

# 设置显示列和显示宽度
pd.set_option('display.max_columns', None)
pd.set_option('display.width', 500)

# 中文设置
plt.rcParams['font.sans-serif'] = 'SimHei'
plt.rcParams['axes.unicode_minus'] = False

# 数据准备
frame = ts.get_k_data('600339', start='2010-01-01', end='2020-10-31')
frame.date = pd.to_datetime(frame['date'])
frame = frame.set_index('date')
print(frame.head())

fig = plt.figure(figsize=(20, 10), dpi=100)
fig.suptitle('股票收盘价趋势图', fontsize=30)

# 操作1：股票收盘价趋势图
plt.subplot2grid((2, 3), (0, 0))
x = frame.index
y = frame.close.values
plt.xlabel('交易日期')
plt.ylabel('收盘价')
plt.title('原始股票收盘价趋势图')
plt.plot(x, y)

# 操作2：趋势图的平滑处理
plt.subplot2grid((2, 3), (0, 1))
y1 = frame['close'].rolling(200, min_periods=2).mean().values
plt.plot(x, y1)
plt.xlabel('交易日期')
plt.ylabel('收盘价')
plt.title('趋势图的平滑处理(rolling方法)')

plt.subplot2grid((2, 3), (0, 2))
y1 = frame['close'].expanding(10).mean().values
plt.plot(x, y1)
plt.xlabel('交易日期')
plt.ylabel('收盘价')
plt.title('趋势图的平滑处理(expanding方法)')

plt.subplot2grid((2, 3), (1, 0))
y1 = frame['close'].ewm(span=30).mean().values
plt.plot(x, y1)
plt.xlabel('交易日期')
plt.ylabel('收盘价')
plt.title('趋势图的平滑处理(指数加权移动方法 - ewm(span = int))')

# 操作3：趋势图的更多时序处理方法
plt.subplot2grid((2, 3), (1, 1))
y1 = frame['close'].pct_change().values
plt.plot(x, y1)
plt.xlabel('交易日期')
plt.ylabel('收盘价')
plt.ylim(-0.3, 0.3)
plt.title('趋势图的平滑处理(百分数变化处理)')

plt.subplot2grid((2, 3), (1, 2))
y1 = frame['close'].rolling(10).mean().apply(lambda i: i * i)
plt.plot(x, y1)
plt.xlabel('交易日期')
plt.ylabel('收盘价')
plt.title('趋势图的平滑处理(自定义函数处理)')

plt.show()
